#include <string>
#include <iostream>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include "cnrt.h"

#ifndef CNRT_CHECK
#define CNRT_CHECK(condition)                   \
    do {                                        \
        condition;                              \
    } while (0)

/* cnrtRet_t ret = condition; \ */
/* CHECK_EQ(ret, CNRT_RET_SUCCESS); \ */
#endif

class ClassifierLauncher {
public:
    explicit ClassifierLauncher();
    void run_network(std::string offmodel, std::string images_file);

    cnrtModel_t model_;
    cnrtFunction_t function_;
    cnrtRuntimeContext_t runtime_ctx_;

    cnrtQueue_t queue_;
    cnrtNotifier_t start_notifier_;
    cnrtNotifier_t end_notifier_;
    int64_t* inputSizeS_, *outputSizeS_;
    int inputNum_, outputNum_;
    void** inputMluPtrS;
    void** outputMluPtrS;
    unsigned int in_n_, in_c_, in_h_, in_w_;
    unsigned int out_n_, out_c_, out_h_, out_w_;
    int out_count_;

    ~ClassifierLauncher() {
        cnrtFreeArray(inputMluPtrS, inputNum_);
        cnrtFreeArray(outputMluPtrS, outputNum_);

        cnrtDestroyFunction(function_);
        cnrtDestroyQueue(queue_);
        cnrtDestroyNotifier(&start_notifier_);
        cnrtDestroyNotifier(&end_notifier_);
        cnrtDestroyRuntimeContext(runtime_ctx_);
        cnrtUnloadModel(model_);
    }
};

ClassifierLauncher::ClassifierLauncher(){}

void ClassifierLauncher::run_network(std::string offmodel, std::string images_file) {
    int dev_id_ = 0;
    cnrtLoadModel(&model_, offmodel.c_str());
    std::string name = "subnet0";

    cnrtDev_t dev;
    CNRT_CHECK(cnrtGetDeviceHandle(&dev, dev_id_));
    CNRT_CHECK(cnrtSetCurrentDevice(dev));
    bool LOG_ON = true;
    if (LOG_ON) {
        std::cout << "Init Classifier for device " << dev_id_ << std::endl;
    }

    CNRT_CHECK(cnrtCreateFunction(&(function_)));
    CNRT_CHECK(cnrtExtractFunction(&(function_), model_, name.c_str()));

    CNRT_CHECK(cnrtCreateRuntimeContext(&runtime_ctx_, function_, nullptr));
    cnrtSetRuntimeContextDeviceId(runtime_ctx_, dev_id_);

    if (cnrtInitRuntimeContext(runtime_ctx_, nullptr) != CNRT_RET_SUCCESS) {
        std::cout << "Failed to init runtime context" << std::endl;
        return;
    }

    CNRT_CHECK(cnrtGetInputDataSize(&inputSizeS_, &inputNum_, function_));
    if (LOG_ON) {
        std::cout << "model input num: " << inputNum_ << " input size " << *inputSizeS_ << std::endl;
    }
    CNRT_CHECK(cnrtGetOutputDataSize(&outputSizeS_, &outputNum_, function_));
    if (LOG_ON) {
        std::cout << "model output num: " << outputNum_ << " output size " << *outputSizeS_ << std::endl;
    }

    inputMluPtrS = (void**)malloc(sizeof(void*) * inputNum_);
    for (int k = 0; k < inputNum_; k++){
        cnrtMalloc(&(inputMluPtrS[k]), inputSizeS_[k]);
    }
    outputMluPtrS = (void**)malloc(sizeof(void*) * outputNum_);
    for (int j = 0; j < outputNum_; j++){
        cnrtMalloc(&(outputMluPtrS[j]), outputSizeS_[j]);
    }

    CNRT_CHECK(cnrtCreateQueue(&queue_));
    CNRT_CHECK(cnrtCreateNotifier(&start_notifier_));
    CNRT_CHECK(cnrtCreateNotifier(&end_notifier_));

    int *dimValues = nullptr;
    int dimNum = 0;
    CNRT_CHECK(cnrtGetInputDataShape(&dimValues, &dimNum, 0, function_));
    in_n_ = dimValues[0];
    in_h_ = dimValues[1];
    in_w_ = dimValues[2];
    in_c_ = dimValues[3];
    free(dimValues);
    if (LOG_ON) {
        std::cout << "model input dimNum: " << dimNum << " N: " << in_n_ << " H: " << in_h_
                  << " W: " << in_w_ << " C: " << in_c_ << std::endl;
    }

    CNRT_CHECK(cnrtGetOutputDataShape(&dimValues, &dimNum, 0, function_));
    out_n_ = dimValues[0];
    out_h_ = dimValues[1];
    out_w_ = 1;
    out_c_ = 1;
    out_count_ = out_n_ * out_h_ * out_w_ * out_c_;
    free(dimValues);
    if (LOG_ON) {
        std::cout << "model output dimNum: " << dimNum << " N: " << out_n_ << " H: " << out_h_
                  << " W: " << out_w_ << " C: " << out_c_ << std::endl;
    }

    int input_size = in_n_ * in_h_ * in_w_ * in_c_;
    void* cpu_data_cast_type = malloc(cnrtDataTypeSize(CNRT_FLOAT16) * input_size);
    cv::Mat img = cv::imread(images_file, cv::IMREAD_GRAYSCALE);
    if (LOG_ON) {
        std::cout << "image w x h x c: " << img.cols << " x " << img.rows
                  << " x " << img.channels() << std::endl;
    }

    float *cpu_data_ = (float *)malloc(input_size * sizeof(float));
    int channels = img.channels();
    for(int k= 0; k < channels; ++k){
        for(int i = 0; i < img.rows; ++i){
            for(int j = 0; j < img.cols; ++j){
                cpu_data_[k*img.rows * img.cols + i * img.cols + j] =
                    img.at<cv::Vec3b>(i, j)[channels - k - 1] / 255.0f;
            }
        }
    }
    for(int i = 0; i < input_size; i++) printf("%f ", cpu_data_[i]);
    
    //cv::Mat img_scale(cv::Size(img.cols, img.cols), CV_32FC1);
    //cv::divide(255.0, img, img_scale, -1);
    //void *cpu_data_ = img_scale.data;
    //float img_data[28*28] = {};
    CNRT_CHECK(cnrtCastDataType(cpu_data_, CNRT_FLOAT32,
                                cpu_data_cast_type, CNRT_FLOAT16, input_size, nullptr));
    CNRT_CHECK(cnrtMemcpy(inputMluPtrS[0], cpu_data_cast_type,
                          inputSizeS_[0], CNRT_MEM_TRANS_DIR_HOST2DEV));
    free(cpu_data_cast_type);

    void* param[inputNum_ + outputNum_];
    for (int j = 0; j < inputNum_; j++) {
        param[j] = inputMluPtrS[j];
    }
    for (int j = 0; j < outputNum_; j++) {
        param[inputNum_ + j] = outputMluPtrS[j];
    }
    
    CNRT_CHECK(cnrtPlaceNotifier(start_notifier_, queue_));
    CNRT_CHECK(cnrtInvokeRuntimeContext(runtime_ctx_, param, queue_, nullptr));
    CNRT_CHECK(cnrtPlaceNotifier(end_notifier_, queue_));

    if (cnrtSyncQueue(queue_) == CNRT_RET_SUCCESS) {
        //cnrtNotifierDuration(start_notifier_, end_notifier_, (float*)(&time_use));
        std::cout << "SyncStream success" << std::endl;
    } else {
        std::cout << "SyncStream error" << std::endl;
    }

    void **output_cpu_data_ = (void **)malloc(sizeof(void **) * outputNum_);
    float **outputCpu = (float**)malloc(sizeof(float*) * outputNum_);
    for (int j = 0; j < outputNum_; j++){
        output_cpu_data_[j] = malloc(sizeof(float) * out_count_);
        CNRT_CHECK(cnrtMemcpy(output_cpu_data_[j], outputMluPtrS[j],
                              outputSizeS_[j], CNRT_MEM_TRANS_DIR_DEV2HOST));
        outputCpu[j] = (float*)malloc(cnrtDataTypeSize(CNRT_FLOAT32) * out_count_);
        CNRT_CHECK(cnrtCastDataType(output_cpu_data_[j], CNRT_FLOAT16,
                                    outputCpu[j], CNRT_FLOAT32, out_count_, nullptr));

        float max_score = -1.0F;
        int max_index = -1;
        for(int i = 0; i < out_count_; i++){
            std::cout << ": " << i << " score " << outputCpu[j][i] << std::endl;
            if(outputCpu[j][i] > max_score){
                max_score = outputCpu[j][i];
                max_index = i;
            }
        }
        std::cout << "class: " << max_index << " score " << max_score << std::endl;
        free(output_cpu_data_[j]);
        free(outputCpu[j]);
    }
    free(output_cpu_data_);
    free(outputCpu);
}

void classification_run(std::string offmodel, std::string images_file) {
    ClassifierLauncher* launcher = new ClassifierLauncher();
    launcher->run_network(offmodel, images_file);
    delete launcher;
}

int main(int argc, char* argv[]) {
    unsigned int real_dev_num;
    cnrtInit(0);
    cnrtGetDeviceCount(&real_dev_num);
    if (real_dev_num == 0) {
        std::cerr << "only have " << real_dev_num << " device(s) " << std::endl;
        cnrtDestroy();
        return -1;
    }

    std::string offmodel = "mnist_deep.cambricon";
    std::string images_file = "9.jpg";
    classification_run(offmodel, images_file);

    cnrtDestroy();
    return 0;
}

/*
  {0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.2,0.62352941,0.99215686,0.62352941,0.19607843,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.18823529,0.93333333,0.98823529,0.98823529,0.98823529,0.92941176,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.21176471,0.89019608,0.99215686,0.98823529,0.9372549,0.91372549,0.98823529,0.22352941,0.02352941,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.03921569,0.23529412,0.87843137,0.98823529,0.99215686,0.98823529,0.79215686,0.32941176,0.98823529,0.99215686,0.47843137,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.63921569,0.98823529,0.98823529,0.98823529,0.99215686,0.98823529,0.98823529,0.37647059,0.74117647,0.99215686,0.65490196,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.2,0.93333333,0.99215686,0.99215686,0.74509804,0.44705882,0.99215686,0.89411765,0.18431373,0.30980392,1.,0.65882353,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.18823529,0.93333333,0.98823529,0.98823529,0.70196078,0.04705882,0.29411765,0.4745098,0.08235294,0.,0.,0.99215686,0.95294118,0.19607843,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.14901961,0.64705882,0.99215686,0.91372549,0.81568627,0.32941176,0.,0.,0.,0.,0.,0.,0.99215686,0.98823529,0.64705882,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.02745098,0.69803922,0.98823529,0.94117647,0.27843137,0.0745098,0.10980392,0.,0.,0.,0.,0.,0.,0.99215686,0.98823529,0.76470588,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.22352941,0.98823529,0.98823529,0.24705882,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.99215686,0.98823529,0.76470588,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.77647059,0.99215686,0.74509804,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,1.,0.99215686,0.76862745,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.29803922,0.96470588,0.98823529,0.43921569,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.99215686,0.98823529,0.58039216,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.33333333,0.98823529,0.90196078,0.09803922,0.,0.,0.,0.,0.,0.,0.,0.,0.02745098,0.52941176,0.99215686,0.72941176,0.04705882,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.33333333,0.98823529,0.8745098,0.,0.,0.,0.,0.,0.,0.,0.,0.02745098,0.51372549,0.98823529,0.88235294,0.27843137,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.33333333,0.98823529,0.56862745,0.,0.,0.,0.,0.,0.,0.,0.18823529,0.64705882,0.98823529,0.67843137,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.3372549,0.99215686,0.88235294,0.,0.,0.,0.,0.,0.,0.44705882,0.93333333,0.99215686,0.63529412,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.33333333,0.98823529,0.97647059,0.57254902,0.18823529,0.11372549,0.33333333,0.69803922,0.88235294,0.99215686,0.8745098,0.65490196,0.21960784,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.33333333,0.98823529,0.98823529,0.98823529,0.89803922,0.84313725,0.98823529,0.98823529,0.98823529,0.76862745,0.50980392,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.10980392,0.78039216,0.98823529,0.98823529,0.99215686,0.98823529,0.98823529,0.91372549,0.56862745,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.09803922,0.50196078,0.98823529,0.99215686,0.98823529,0.55294118,0.14509804,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.}
*/
